module geog_data_mod

  use spark

  implicit none

  type geog_domain_array_type
    ! 1: quadrature point
    ! 2: x-axis
    ! 3: y-axis
    ! 4: z-axis
    ! 5: variable
    
    !                                1 2 3
    real(r_kind), pointer, dimension(  :,:) :: zs      ! Surface height on cell average
    real(r_kind), pointer, dimension(  :,:) :: zsC     ! Surface height on cell center point
    real(r_kind), pointer, dimension(:,:,:) :: zsL     ! Surface height on left points
    real(r_kind), pointer, dimension(:,:,:) :: zsR     ! Surface height on right points
    real(r_kind), pointer, dimension(:,:,:) :: zsB     ! Surface height on bottom points
    real(r_kind), pointer, dimension(:,:,:) :: zsT     ! Surface height on top points
    real(r_kind), pointer, dimension(:,:,:) :: zsQ     ! Surface height on quadrature points
    
    real(r_kind), pointer, dimension(  :,:) :: dzsdx
    real(r_kind), pointer, dimension(  :,:) :: dzsdxC
    real(r_kind), pointer, dimension(:,:,:) :: dzsdxL
    real(r_kind), pointer, dimension(:,:,:) :: dzsdxR
    real(r_kind), pointer, dimension(:,:,:) :: dzsdxB
    real(r_kind), pointer, dimension(:,:,:) :: dzsdxT
    real(r_kind), pointer, dimension(:,:,:) :: dzsdxQ
    
    real(r_kind), pointer, dimension(  :,:) :: dzsdy
    real(r_kind), pointer, dimension(  :,:) :: dzsdyC
    real(r_kind), pointer, dimension(:,:,:) :: dzsdyL
    real(r_kind), pointer, dimension(:,:,:) :: dzsdyR
    real(r_kind), pointer, dimension(:,:,:) :: dzsdyB
    real(r_kind), pointer, dimension(:,:,:) :: dzsdyT
    real(r_kind), pointer, dimension(:,:,:) :: dzsdyQ
    
    real(r_kind), pointer, dimension(:,:) :: landfrac
    real(r_kind), pointer, dimension(:,:) :: slmsk
  end type geog_domain_array_type

  type geog_data_type
    logical :: initialized = .false.
    type(cubed_sphere_mesh_type), pointer :: mesh => null()
    type(cubed_sphere_array_type) array
    type(geog_domain_array_type), allocatable :: domain(:)
  contains
    procedure :: init   => geog_data_init
    procedure :: clear  => geog_data_clear
    final :: geog_data_final
  end type geog_data_type

contains

  subroutine geog_data_init(this, mesh)

    class(geog_data_type), intent(inout) :: this
    type(cubed_sphere_mesh_type), intent(in), target :: mesh

    integer idom

    call this%clear()

    this%mesh => mesh

    call this%array%init(mesh)
    call this%array%add_var('zs', 'Surface height', 'm', loc='C:CQ:LQ:RQ:TQ:BQ:CA', with_halo='T:T:T:T:T:T:T', fill_halo='F:F:F:F:F:F:F', output='T:T:T:T:T:T:T', only_2d=.true.)
    
    call this%array%add_var('dzsdx', 'Surface height derivative in x-dir', 'm/rad', loc='C:CQ:LQ:RQ:TQ:BQ:CA', with_halo='T:T:T:T:T:T:T', fill_halo='F:F:F:F:F:F:F', output='T:T:T:T:T:T:T', only_2d=.true.)
    call this%array%add_var('dzsdy', 'Surface height derivative in y-dir', 'm/rad', loc='C:CQ:LQ:RQ:TQ:BQ:CA', with_halo='T:T:T:T:T:T:T', fill_halo='F:F:F:F:F:F:F', output='T:T:T:T:T:T:T', only_2d=.true.)

    call this%array%add_var('landfrac', 'Land Fraction', '-', loc='CA', with_halo='T', fill_halo='F', output='T', only_2d=.true.)
    call this%array%add_var('slmsk'   , 'Sea land mask', '-', loc='CA', with_halo='T', fill_halo='F', output='T', only_2d=.true.)
    call this%array%allocate_arrays()

    allocate(this%domain(mesh%ndom))
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%initialized) then
        call this%array%get_array(this%domain(idom)%zs      , idom=idom, var_name='zs'      , loc='CA')
        call this%array%get_array(this%domain(idom)%zsC     , idom=idom, var_name='zs'      , loc='C' )
        call this%array%get_array(this%domain(idom)%zsL     , idom=idom, var_name='zs'      , loc='LQ')
        call this%array%get_array(this%domain(idom)%zsR     , idom=idom, var_name='zs'      , loc='RQ')
        call this%array%get_array(this%domain(idom)%zsB     , idom=idom, var_name='zs'      , loc='BQ')
        call this%array%get_array(this%domain(idom)%zsT     , idom=idom, var_name='zs'      , loc='TQ')
        call this%array%get_array(this%domain(idom)%zsQ     , idom=idom, var_name='zs'      , loc='CQ')
        
        call this%array%get_array(this%domain(idom)%dzsdx      , idom=idom, var_name='dzsdx'      , loc='CA')
        call this%array%get_array(this%domain(idom)%dzsdxC     , idom=idom, var_name='dzsdx'      , loc='C' )
        call this%array%get_array(this%domain(idom)%dzsdxL     , idom=idom, var_name='dzsdx'      , loc='LQ')
        call this%array%get_array(this%domain(idom)%dzsdxR     , idom=idom, var_name='dzsdx'      , loc='RQ')
        call this%array%get_array(this%domain(idom)%dzsdxB     , idom=idom, var_name='dzsdx'      , loc='BQ')
        call this%array%get_array(this%domain(idom)%dzsdxT     , idom=idom, var_name='dzsdx'      , loc='TQ')
        call this%array%get_array(this%domain(idom)%dzsdxQ     , idom=idom, var_name='dzsdx'      , loc='CQ')
        
        call this%array%get_array(this%domain(idom)%dzsdy      , idom=idom, var_name='dzsdy'      , loc='CA')
        call this%array%get_array(this%domain(idom)%dzsdyC     , idom=idom, var_name='dzsdy'      , loc='C' )
        call this%array%get_array(this%domain(idom)%dzsdyL     , idom=idom, var_name='dzsdy'      , loc='LQ')
        call this%array%get_array(this%domain(idom)%dzsdyR     , idom=idom, var_name='dzsdy'      , loc='RQ')
        call this%array%get_array(this%domain(idom)%dzsdyB     , idom=idom, var_name='dzsdy'      , loc='BQ')
        call this%array%get_array(this%domain(idom)%dzsdyT     , idom=idom, var_name='dzsdy'      , loc='TQ')
        call this%array%get_array(this%domain(idom)%dzsdyQ     , idom=idom, var_name='dzsdy'      , loc='CQ')
        
        call this%array%get_array(this%domain(idom)%landfrac, idom=idom, var_name='landfrac', loc='CA')
        call this%array%get_array(this%domain(idom)%slmsk   , idom=idom, var_name='slmsk'   , loc='CA')
      end if
    end do

    this%initialized = .true.

  end subroutine geog_data_init

  subroutine geog_data_clear(this)

    class(geog_data_type), intent(inout) :: this

    nullify(this%mesh)

    call this%array%clear()

    if (allocated(this%domain)) deallocate(this%domain)

    this%initialized = .false.

  end subroutine geog_data_clear

  subroutine geog_data_final(this)

    type(geog_data_type), intent(inout) :: this

    call this%clear()

  end subroutine geog_data_final

end module geog_data_mod
